#include "stdafx.h"

#include "HeapMgr.h"
#include "Txn.h"
#include "TxnMgr.h"
#include <tinyxml.h>
#include <algorithm>

#define PAGE_SIZE  4096
#define BLOCK_SIZE (PAGE_SIZE * 16)

HeapMgr::HeapMgr(void)
{
}

HeapMgr::~HeapMgr(void)
{
	Clear();
}

bool HeapMgr::Clear()
{
	bool bOK = false; 

	std::vector<void*>::const_iterator itEnd = m_vPageAllocated.end();
	std::vector<void*>::const_iterator it;
	for(it = m_vPageAllocated.begin(); it != itEnd; ++it)
		::VirtualFree(*it, 0, MEM_RELEASE);

	m_vPageAllocated.clear();
	m_mMemoryUsed.clear();
	m_mMemoryUnused.clear();

	return bOK;
}


HeapMgr* HeapMgr::Get()
{
	static HeapMgr heapMgr;
	return &heapMgr;
}

// Exception filter
// Backup the target page
LONG	HeapMgr::ExceptionFilter(PEXCEPTION_POINTERS pVEHExp)
{
	// Must happen in a transaction
	Txn* pTxn = TxnMgr::GetActiveTxn();
	if(!pTxn)
		return EXCEPTION_CONTINUE_SEARCH;

	LPEXCEPTION_POINTERS e = static_cast<LPEXCEPTION_POINTERS>(pVEHExp);

	// Ignore if the exception is not EXCEPTION_ACCESS_VIOLATION
	if (e->ExceptionRecord->ExceptionCode != EXCEPTION_ACCESS_VIOLATION)
		return EXCEPTION_CONTINUE_SEARCH;

	// The first element of the ExceptionInformation[] array contains a read-write flag 
	// that indicates the type of operation that caused the access violation. 
	// If this value is 0, the thread attempted to read the inaccessible data. 
	// If this value is 1, the thread attempted to write to an inaccessible address. 
	// The second array element specifies the virtual address of the inaccessible data. 
	bool	bWriting = (e->ExceptionRecord->ExceptionInformation[0] != 0);
	void*	pAddr	 = (void*)e->ExceptionRecord->ExceptionInformation[1];
	void*	pPage	 = GetPageAddress(pAddr);//(void*)((ULONG_PTR)pAddr & ~(PAGE_SIZE - 1));

	// Ignore if reading data 
	if (!bWriting)
		return EXCEPTION_CONTINUE_SEARCH;

	// Ignore if this memory is not controlled by this manager
	if (!IsStorage(pPage))
		return EXCEPTION_CONTINUE_SEARCH;

	// Backup this page
	void* pBackup = BackupAPage(pPage);
	if (!pBackup)
		return EXCEPTION_CONTINUE_SEARCH;

	BOOL bOK = pTxn->SetMap(pPage, pBackup);
	if(!bOK)
		return EXCEPTION_CONTINUE_SEARCH;

	// Make the original memory can be written into
	bOK = UnlockStorage(pPage, PAGE_SIZE);
	if(!bOK)
		return EXCEPTION_CONTINUE_SEARCH;

	return EXCEPTION_CONTINUE_EXECUTION;
}

// Get one page's size
size_t HeapMgr::GetPageSize() const
{
	return PAGE_SIZE;
}

// Get page's address
// Remove the extra from pointer address
void* HeapMgr::GetPageAddress(void* pAddress) const
{
	return (void*)((ULONG_PTR)pAddress & ~(PAGE_SIZE - 1));
}

// Allocate memory 
void* HeapMgr::AllocStorage(size_t nSize)
{
	void* pAddress = NULL;

	// Try to allocate the memory from the unused block
	std::map<void*, size_t>::const_iterator itEnd = m_mMemoryUnused.end();
	std::map<void*, size_t>::iterator it;
	for(it = m_mMemoryUnused.begin(); it != itEnd; ++it)
	{
		if(it->second == nSize)
		{
			pAddress = it->first;
			m_mMemoryUsed[pAddress] = nSize;
			m_mMemoryUnused.erase(it);
			break;
		}
		else if(it->second > nSize)
		{
			pAddress = it->first;
			m_mMemoryUsed[pAddress] = nSize;
			size_t nLeft = it->second - nSize;
			m_mMemoryUnused.erase(it);
			m_mMemoryUnused[(void*)((char*)(pAddress)+nSize)] = nLeft;
			break;
		}
	}

	// If no available memory in cache, VirtualAlloc it
	if(!pAddress)
	{
		// Get pages that will be allocated
		size_t iPage = nSize/PAGE_SIZE;
		if(nSize%PAGE_SIZE != 0)
			iPage++;

		// Get the real size that will be allocated
		size_t nRealSize = iPage*PAGE_SIZE;

		// Allocate memory
		// It needs to be read-write as we need to initialize the memory out of
		// a transaction, like a vector.
		pAddress = ::VirtualAlloc(NULL, nRealSize, MEM_COMMIT, PAGE_READWRITE);

		// Record the memory's information
		if(pAddress)
		{
			// Store the allocated page for release in the future
			m_vPageAllocated.push_back(pAddress);

			// Record the used memory's information
			m_mMemoryUsed[pAddress] = nSize;

			// Record the unused memory's information
			if(nRealSize > nSize)
				m_mMemoryUnused[(void*)((char*)(pAddress)+nSize)] = nRealSize-nSize;
		}
	}


	return pAddress;
}

// Deallocate memory of storage
BOOL HeapMgr::FreeStorage(void* pAddress)
{
	// Set the memory to unused state.
	std::map<void*, size_t>::const_iterator itEnd = m_mMemoryUsed.end();
	std::map<void*, size_t>::iterator it;
	for(it = m_mMemoryUsed.begin(); it != itEnd; ++it)
	{
		if (it->first == pAddress)
		{
			m_mMemoryUnused[it->first] = it->second;
			m_mMemoryUsed.erase(it);

			return TRUE;
		}
	}

	TXN_ASSERT(FALSE && "Hey, I am not your manager!!!");

	return FALSE;
}

// Get memory information
BOOL HeapMgr::GetStorage(void* pAddress, void** p, size_t* pnSize) const
{
	std::map<void*, size_t>::const_iterator itEnd = m_mMemoryUsed.end();
	std::map<void*, size_t>::const_iterator it;
	for(it = m_mMemoryUsed.begin(); it != itEnd; ++it)
	{
		if(it->first <= pAddress && (char*)it->first+it->second > pAddress)
		{
			if(p)
				*p = it->first;
			if(pnSize)
				*pnSize = it->second;
			return TRUE;
		}
	}

	return FALSE;
}

// Is the memory controlled by this manager
BOOL HeapMgr::IsStorage(void* pAddress) const
{
	return (GetStorage(pAddress)!=NULL);
}

// Lock all storage memory
BOOL HeapMgr::LockAll() const
{
	BOOL bOK = FALSE;

	std::map<void*, size_t>::const_iterator itEnd = m_mMemoryUsed.end();
	std::map<void*, size_t>::const_iterator it;
	for(it = m_mMemoryUsed.begin(); it != itEnd; ++it)
	{
		bOK = LockStorage(it->first, it->second);
		TXN_ASSERT(bOK);
	}

	return bOK;
}

// Unlock all storage memory
BOOL HeapMgr::UnlockAll() const
{
	BOOL bOK = FALSE;

	std::map<void*, size_t>::const_iterator itEnd = m_mMemoryUsed.end();
	std::map<void*, size_t>::const_iterator it;
	for(it = m_mMemoryUsed.begin(); it != itEnd; ++it)
	{
		bOK = UnlockStorage(it->first, it->second);
		TXN_ASSERT(bOK);
	}

	return bOK;
}

// Whether the page is locked
BOOL HeapMgr::IsLocked(void* pAddress) const
{
	// Get page's address (pAddress is in that page)
	void* pPage = GetPageAddress(pAddress);

	MEMORY_BASIC_INFORMATION info;
	::VirtualQuery(pPage, &info, PAGE_SIZE);
	if(info.Protect == PAGE_READONLY)
		return TRUE;

	return FALSE;
}

// Lock memory
BOOL HeapMgr::LockStorage(void* pAddress, size_t nSize) const
{
	// Get page's address (pAddress is in that page)
	void* pPage = GetPageAddress(pAddress);

	// Get page number
	nSize = (size_t)pAddress + nSize - (size_t)pPage;
	size_t nPage = size_t(nSize/PAGE_SIZE);
	if(nSize%PAGE_SIZE != 0)
		nPage++;

	DWORD dwOld = 0;
	// Make the backup memory can't be written into
	return ::VirtualProtect(pPage, PAGE_SIZE*nPage, PAGE_READONLY, &dwOld);
}

// Unlock memory
BOOL HeapMgr::UnlockStorage(void* pAddress, size_t nSize) const
{
	// Get page's address (pAddress is in that page)
	void* pPage = GetPageAddress(pAddress);

	// Get page number
	nSize = (size_t)pAddress + nSize - (size_t)pPage;
	size_t nPage = size_t(nSize/PAGE_SIZE);
	if(nSize%PAGE_SIZE != 0)
		nPage++;

	DWORD dwOld = 0;
	// Make the backup memory can be written into
	return ::VirtualProtect(pPage, PAGE_SIZE*nPage, PAGE_READWRITE, &dwOld);
}

// Backup a page memory
// return the base address of the "to" memory
void* HeapMgr::BackupAPage(void* pAddress)
{
	TXN_ASSERT(IsStorage(pAddress));

	// Get the page used to store the memory
	void* pPageStorage = GetPageAddress(pAddress);
	if(!pPageStorage)
		return NULL;

	// Allocate memory for backup - backup to crt heap, not txn heap
	void* pPageBackup = new char [PAGE_SIZE];

	// Backup a page that includes pAddress
	memcpy(pPageBackup, pPageStorage, PAGE_SIZE);

	// The backup memory
	return pPageBackup;
}

struct Block
{
	void* addr;
	size_t size;
	int type; // 1 = used, 0 = unused
};

static bool BlockCompare(const Block& lhs, const Block& rhs)
{
	return lhs.addr < rhs.addr;
}

void HeapMgr::Dump()
{

	// Build the Page->Blocks map to reflect memory status
	typedef std::map<void*, std::vector<Block> >  PageBlocksMap;

	// All pages to map
	PageBlocksMap mapPageBlocks;
	std::vector<void*>::const_iterator itp;
	for (itp = m_vPageAllocated.begin(); itp != m_vPageAllocated.end(); itp++)
	{
		mapPageBlocks[*itp] = std::vector<Block> ();
	}


	// Walk through the used memory blocks
	std::map<void*, size_t>::const_iterator it;
	for (it = m_mMemoryUsed.begin(); it != m_mMemoryUsed.end(); it++)
	{
		void* pPage = GetPageAddress(it->first);
		Block b = {it->first, it->second, 1};
		mapPageBlocks[pPage].push_back(b);
	}

	// Walk through the unused memory blocks
	for (it = m_mMemoryUnused.begin(); it != m_mMemoryUnused.end(); it++)
	{
		void* pPage = GetPageAddress(it->first);
		Block b = {it->first, it->second, 0};
		mapPageBlocks[pPage].push_back(b);
	}


	// Output to XML
	// <Pages>
	//		<Page address=00000000>
	//			<Block address=11111111 used=1>10byte</Block>
	//		</Page>
	// </Pages>
	//
	TiXmlDocument doc;
	TiXmlDeclaration * decl = new TiXmlDeclaration( "1.0", "", "" );
	doc.LinkEndChild( decl );

	// "Pages"
	TiXmlElement* pAllocatedPages = new TiXmlElement( "Pages" );
	doc.LinkEndChild(pAllocatedPages);


	PageBlocksMap::iterator itpb;
	for (itpb = mapPageBlocks.begin(); itpb != mapPageBlocks.end(); itpb++)
	{
		// Page
		TiXmlElement* pPage = new TiXmlElement( "Page" );
		pAllocatedPages->LinkEndChild(pPage);

		char addr[20];
		sprintf_s(addr, "%p", itpb->first);
		pPage->SetAttribute("address", addr);

		// Blocks
		std::vector<Block>& blocks = itpb->second;
		std::sort(blocks.begin(), blocks.end(), BlockCompare);

		for (std::vector<Block>::const_iterator itb = blocks.begin(); itb != blocks.end(); itb++)
		{
			// Block
			TiXmlElement* pBlock = new TiXmlElement( "Block" );
			pPage->LinkEndChild(pBlock);

			// address
			char addr[20];
			sprintf_s(addr, "%p", itb->addr);
			pBlock->SetAttribute("address", addr);

			// type
			char type[20];
			sprintf_s(type, "%d", itb->type);
			pBlock->SetAttribute("used", type);

			// size
			char size[20];
			sprintf_s(size, "%5d", itb->size);
			TiXmlText * text = new TiXmlText( size);
			pBlock->LinkEndChild( text );
		}

	}


	doc.SaveFile("C:\\HeapMgr.txt");
}

LONG	NTAPI TxnExceptionHandler(PEXCEPTION_POINTERS pVEHExp)
{
	return HeapMgr::Get()->ExceptionFilter(pVEHExp);
}

class VEHAutoInstaller
{
public:
	VEHAutoInstaller()
	{
		::AddVectoredExceptionHandler(1, &TxnExceptionHandler);
	}
	~VEHAutoInstaller()
	{
		::RemoveVectoredExceptionHandler(&TxnExceptionHandler);
	}
};

// Install the TxnExceptionHandler at application startup
static VEHAutoInstaller installer;